'''
A quick and dirty VTK based viewer for volume data.
'''


import numpy, wx, vtk, defacing ,threading,datetime
import os
try:       
    from vtkImageImportFromArray import vtkImageImportFromArray
except ImportError:
    print 'You need to get the updated vtkImageImportFromArray from: "http://www.siafoo.net/snippet/313" '
    exit(1)


from threading import Thread

def run(data,thres):
        
    ren = vtk.vtkRenderer()
    renWin = vtk.vtkRenderWindow()
    renWin.SetWindowName(str(thres))
    renWin.AddRenderer(ren)
    iren = vtk.vtkRenderWindowInteractor()
    iren.SetRenderWindow(renWin)
    img,imax,imin = LoadVolumeData(data,thres)

    pix_diag = 5.0


    volMapper = vtk.vtkVolumeRayCastMapper()
    compositeFunction = vtk.vtkVolumeRayCastCompositeFunction()
    compositeFunction.SetCompositeMethodToInterpolateFirst()
    volMapper.SetVolumeRayCastFunction(compositeFunction)

    volMapper.SetSampleDistance(pix_diag / 5.0)
    #er.CroppingOn()
    volMapper.SetCroppingRegionPlanes(00.0,10000.0,0.0,100.0,0.0,10000.0)
    volMapper.SetCroppingRegionFlagsToSubVolume()
    volMapper.SetInputConnection(img.GetOutputPort())

    # Transfer Functions
    opacity_tf = vtk.vtkPiecewiseFunction()
    color_tf = vtk.vtkColorTransferFunction()

    # The property describes how the data will look

    color_tf.AddRGBPoint(imax, 1,1,1)
    color_tf.AddRGBPoint(0, 0.2,0.2,0.2)
    opacity_tf.AddPoint(imin,0)
    opacity_tf.AddPoint(imax,1)

    volProperty = vtk.vtkVolumeProperty()
    volProperty.SetShade(False)
    volProperty.SetColor(color_tf)
    volProperty.SetScalarOpacity(opacity_tf)
    volProperty.ShadeOn()
    volProperty.SetInterpolationTypeToLinear()
    volProperty.SetScalarOpacityUnitDistance(pix_diag)

    vol = vtk.vtkVolume()
    vol.SetMapper(volMapper)
    vol.SetProperty(volProperty)
        
    ren.AddVolume(vol)
        
    contours= defacing.ExtractSurface(img)
        
    
    normals = defacing.ExtractLargest(contours)
    
    

    normals=defacing.SmoothSurface(normals,200)

    contMapper = vtk.vtkPolyDataMapper()
    contMapper.SetInputConnection(normals.GetOutputPort())
    contMapper.ScalarVisibilityOff()
      
        
    contActor = vtk.vtkActor()
    contActor.SetMapper(contMapper)

    contActor.GetProperty().SetDiffuseColor(1, .49, .25)
    #contActor.GetProperty().SetSpecular(.3)
    contActor.GetProperty().SetSpecularPower(20)


    # ren.AddActor(contActor)

    
    iren.Initialize()
    renWin.Render()
    iren.Start()

def save_decimated(data,thres,level,fname):
        
    ren = vtk.vtkRenderer()
    renWin = vtk.vtkRenderWindow()
    renWin.SetWindowName(str(thres))
    renWin.AddRenderer(ren)
    iren = vtk.vtkRenderWindowInteractor()
    iren.SetRenderWindow(renWin)
    img,imax,imin = LoadVolumeData(data,thres)

    pix_diag = 5.0


    volMapper = vtk.vtkVolumeRayCastMapper()
    compositeFunction = vtk.vtkVolumeRayCastCompositeFunction()
    compositeFunction.SetCompositeMethodToInterpolateFirst()
    volMapper.SetVolumeRayCastFunction(compositeFunction)

    volMapper.SetSampleDistance(pix_diag / 5.0)
    #er.CroppingOn()
    volMapper.SetCroppingRegionPlanes(00.0,10000.0,0.0,100.0,0.0,10000.0)
    volMapper.SetCroppingRegionFlagsToSubVolume()
    volMapper.SetInputConnection(img.GetOutputPort())

    # Transfer Functions
    opacity_tf = vtk.vtkPiecewiseFunction()
    color_tf = vtk.vtkColorTransferFunction()

    # The property describes how the data will look

    color_tf.AddRGBPoint(imax, 1,1,1)
    color_tf.AddRGBPoint(0, 0.2,0.2,0.2)
    opacity_tf.AddPoint(imin,0)
    opacity_tf.AddPoint(imax,1)

    volProperty = vtk.vtkVolumeProperty()
    volProperty.SetShade(False)
    volProperty.SetColor(color_tf)
    volProperty.SetScalarOpacity(opacity_tf)
    volProperty.ShadeOn()
    volProperty.SetInterpolationTypeToLinear()
    volProperty.SetScalarOpacityUnitDistance(pix_diag)

    vol = vtk.vtkVolume()
    vol.SetMapper(volMapper)
    vol.SetProperty(volProperty)
        
    #ren.AddVolume(vol)
        
    contours= defacing.ExtractSurface(img)
        
    
    normals = defacing.ExtractLargest(contours)	
    

    normals=defacing.SmoothSurface(normals,200)
    decimated = defacing.Decimate(normals,level)
    contMapper = vtk.vtkPolyDataMapper()
    contMapper.SetInputConnection(decimated.GetOutputPort())
    contMapper.ScalarVisibilityOff()
      
        
    contActor = vtk.vtkActor()
    contActor.SetMapper(contMapper)

    contActor.GetProperty().SetDiffuseColor(1, .49, .25)
    #contActor.GetProperty().SetSpecular(.3)
    contActor.GetProperty().SetSpecularPower(20)

    PLYwriter = vtk.vtkSTLWriter()
    tmp = os.path.splitext(fname)
        
    PLYwriter.SetFileName('%s_%03d_%03d.stl'%(tmp[0],int(thres),int(100*level)));
    PLYwriter.SetInputConnection(decimated.GetOutputPort());
    PLYwriter.Write();

    ren.AddActor(contActor)

    iren.Initialize()
    renWin.Render()
    iren.Start()
 

def LoadVolumeData( data,thres):
 

    img_data = numpy.cast['uint8'](data)               
    img = vtkImageImportFromArray()

    #thread.start_new_thread(defacing.FrontOrthoRayCast, (img_data,thres,data_set))
  
    #  img_data = defacing.FullOrthoCast(img_data,thres,False,True);

    #print "starting"
    img.SetArray(img_data) #Edited Costas 6/29/2011d	
   
    return img,img_data.max(),img_data.min()

def save_show( data,thres,level,fname):

    print 'fullortho'
    data = defacing.FullOrthoCast(data,thres,False,True)   
    img_data = numpy.cast['uint8'](data)       
    print 'show save'
    img = vtkImageImportFromArray()


    img.SetArray(img_data) #Edited Costas 6/29/2011d	
    Thread(target=save_decimated, args=(data,thres,level,fname)).start()
    return img,img_data.max(),img_data.min()
    

def show(data,thres):
  #  thread.start_new_thread(run, (data,thres))
   Thread(target=run, args=(data,thres)).start()
   

